# https://www.kaggle.com/competitions/digit-recognizer
import pandas as pd
import numpy as np
import joblib 
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split

# 1. 读取数据
train = pd.read_csv('./train.csv')  # 包含 label 和像素特征
test = pd.read_csv('./test.csv')    # 只有像素特征

# 2. 数据预处理
X = train.drop('label', axis=1).values  # 特征
y = train['label'].values               # 标签

# 可以进行归一化（提升模型效果）
X = X / 255.0
test = test.values / 255.0


# 3. 划分训练集和验证集
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)

# 4. 模型训练（这里用随机森林，也可以换成 SVM、KNN、逻辑回归等）
clf = RandomForestClassifier(n_estimators=100, random_state=42)
clf.fit(X_train, y_train)

# 5. 验证集评估
y_pred_val = clf.predict(X_val)
print('Validation Accuracy:', accuracy_score(y_val, y_pred_val))

# 6. 对测试集进行预测
y_pred_test = clf.predict(test)

# 7. 保存训练模型 
joblib.dump(clf, './digit_recognizer_model.pkl')
print('Model saved: ./digit_recognizer_model.pkl')

# 8. 生成提交文件
submission = pd.DataFrame({
    'ImageId': np.arange(1, len(y_pred_test) + 1),
    'Label': y_pred_test
})

submission.to_csv('submission.csv', index=False)
print('Submission file generated: submission.csv')